import json
import pandas as pd
import matplotlib.pyplot as plt

TEXT_LEN = 11

as_dataframe = pd.read_csv("../../src/extracted/DESED_Audioset_label_mapping.csv")
# remove desed_labels only contains speech
non_single_speech = as_dataframe['desed_labels'].apply(lambda x: (set(["speech", "other"]) != set(x.split(","))) and (set(["speech"]) != set(x.split(","))))
print(f"Number of Total: {len(as_dataframe)}, non-single speech labels: {non_single_speech.sum()}") 
non_speech_df = as_dataframe[non_single_speech]
# remove desed files 
# load all desed filenames
desed_filenames = "../../src/DESED_all_filenames.json"
with open(desed_filenames, "r") as f:
    desed_filenames = set(json.load(f))
non_desed = non_speech_df['desed_filename'].apply(lambda x: x not in desed_filenames)
print(f"Number of non-desed files: {non_desed.sum()} out of {len(non_speech_df)}")
selected_df = non_speech_df[non_desed]

as_strong_dataframe = pd.read_csv("../../src/extracted/DESED_AudiosetStrong_label_mapping_non_single_speech.csv")
as_now = selected_df["# YTID"]
as_strong = [x[:TEXT_LEN] for x in as_strong_dataframe['segment_id'].values]
print(as_strong[:10])
# assert all as_strong in same length
l = len(as_strong[0])
for x in as_strong:
    assert len(x) == l, f"Length mismatch: {x} has length {len(x)}, expected {l}"
as_now_dup = as_now.apply(lambda x: x in as_strong)
selected_df = selected_df[~as_now_dup]
# print(selected_df[as_prev_dup]["segment_id"].values[:10])
# non_prev = as_prev_speech.apply(lambda x: "Y" + x.split("_")[0] in non_speech_df['segment_id'].values)
print(f"Number of non-prev files: {as_now_dup.sum()} out of {len(as_now_dup)}, final: {len(selected_df)}")

selected_df.to_csv("../../src/extracted/DESED_Audioset_label_mapping_non_single_speech.csv", index=False)
# plot the histogram of the number of labels
desed_labels = ",".join(selected_df['desed_labels'].values).replace("other,", "").replace(",other", "")
desed_labels = desed_labels.split(",")
desed_labels = pd.Series(desed_labels)
as_labels = desed_labels.value_counts()
plt.title("Number of labels in non-single speech labels")
plt.xlabel("Labels")
plt.ylabel("Count") 

# convert to percentage dictionary
desed_labels_percentage = desed_labels.value_counts(normalize=True).to_dict()
print("Percentage of labels in non-single speech labels:")
for label, percentage in desed_labels_percentage.items():
    print(f"{label}: {percentage:.2%}")

simplified_label = {
    "speech": "speech",
    "dog": "dog",
    "cat": "cat",
    "alarm_bell_ringing": "alarm",
    "dishes": "dishes",
    "frying": "frying",
    "blender": "blender",
    "running_water": "water",
    "vacuum_cleaner": "vacuum",
    "electric_shaver_toothbrush": "shaver"
}

weak_label_counts = {
    "speech": 550,
    "dog": 214,
    "cat": 173,
    "alarm_bell_ringing": 205,
    "dishes": 184,
    "frying": 171,
    "blender": 134,
    "running_water": 343,
    "vacuum_cleaner": 167,
    "electric_shaver_toothbrush": 103
}

weak_label_counts = {simplified_label.get(k, k): v for k, v in weak_label_counts.items()}
as_labels = as_labels.rename(simplified_label)

# plot the weak label counts
plt.figure(figsize=(10, 6))
# Convert weak_label_counts to percentages
weak_total = sum(weak_label_counts.values())
weak_label_percentages = {k: (v/weak_total)*100 for k, v in weak_label_counts.items()}

# Convert as_labels to percentages
as_labels_percentages = (as_labels / as_labels.sum()) * 100

# Get all unique labels for consistent x-axis
all_labels = list(set(list(weak_label_percentages.keys()) + list(as_labels_percentages.index)))
all_labels.sort()
# Prepare data for bar plot
weak_values = [weak_label_percentages.get(label, 0) for label in all_labels]
desed_values = [as_labels_percentages.get(label, 0) for label in all_labels]

x = range(len(all_labels))
width = 0.35

plt.bar([i - width/2 for i in x], weak_values, width, label='Weak labels', alpha=0.7)
plt.bar([i + width/2 for i in x], desed_values, width, label='AS extracted labels', alpha=0.7)
plt.xticks(x, all_labels)
plt.title("Number of labels in non-single speech labels")
plt.xlabel("Labels")
plt.ylabel("Percentage (%)") 
plt.xticks(rotation=90)
plt.tight_layout()
plt.legend()
plt.savefig("../../src/extracted/DESED_Audioset_label_mapping_non_single_speech_histogram.png")